# Copyright (c) 2025 Huawei Technologies Co., Ltd.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import sys

import yaml

# Dictionary defining legacy arguments and their mappings to new keys, along with metadata.
LEGACY_ARGS = {
    "adapter_models": {
        "new_key": "adapter_name_or_path",
        "help": "Path to the adapter weight.",
        "type": str,
        "default": None,
    },
    "yaml_path": {
        "new_key": None,
        "help": "Path to the yaml file.",
        "type": str,
        "default": None,
    },
    "infer_dtype": {
        "new_key": None,
        "help": "Specify inference data type. If 'bfloat16', then bf16=True; if 'float16', then fp16=True; otherwise both are False.",
        "type": str,
        "default": None,
    },
    "hub_token": {
        "new_key": "token",
        "help": "Token to download model from hub.",
        "type": str,
        "default": None,
    },
    "model": {
        "new_key": "model_name_or_path",
        "help": "Model for eval.",
        "type": str,
        "default": None,
    },
    "npu_device_ids": {
        "new_key": "device",
        "help": "npu ids allocated to the model instance.",
        "type": str,
        "default": None,
    },
}


def _add_legacy_args(parser):
    """Adds legacy arguments to the argument parser."""

    group = parser.add_argument_group(title="legacy")
    for legacy_arg, config in LEGACY_ARGS.items():
        group.add_argument(
            f"--{legacy_arg}",
            type=config.get("type", str),
            default=config.get("default", None),
            help=config["help"],
        )
    return parser


def _migrate_legacy_args(parser, params_dict, unknown_args=None):
    """
    Replaces legacy argument keys with their new equivalents where applicable. And remove legacy arguments
    """
    if params_dict["yaml_path"] is not None:
        yaml_path = params_dict["yaml_path"]

        with open(yaml_path, "r") as f:
            yaml_args = yaml.safe_load(f)

        defined_params = {action.dest: action for action in parser._actions}
        known_args = {}
        extra_args = {}

        cli_args_dict = {
            sys.argv[i].lstrip("--"): sys.argv[i + 1]
            for i in range(1, len(sys.argv) - 1)
            if sys.argv[i].startswith("--") and sys.argv[i] != "--yaml_path"
        }

        for key, value in yaml_args.items():
            if key in defined_params:
                if key in cli_args_dict:
                    continue
                known_args[key] = value
            else:
                extra_args[key] = value

        if extra_args:
            raise ValueError(f"Currently {extra_args} is not supported.")

        params_dict.update(known_args)

    # unknown_args ["chat","AI_Connect/Qwen2-0.5B"] or ["chat","xx/xx/Qwen2-0.5B"], which is applicable when run with "openmind-cli chat AI_Connect/Qwen2-0.5B"
    if not unknown_args:
        pass
    elif unknown_args and len(unknown_args) == 2:
        params_dict["model_name_or_path"] = unknown_args[-1]
    else:
        raise ValueError(f"Please check the arguments {unknown_args[2::]}.")

    legacy_infer_dtype = params_dict.get("infer_dtype")
    if legacy_infer_dtype is not None:
        dtype_lower = legacy_infer_dtype.lower()
        if dtype_lower == "bfloat16":
            params_dict["bf16"] = True
            params_dict["fp16"] = False
        elif dtype_lower == "float16":
            params_dict["fp16"] = True
            params_dict["bf16"] = False
        elif dtype_lower == "float32":
            params_dict["bf16"] = False
            params_dict["fp16"] = False
        else:
            raise ValueError(f"Unsupported dtype: {dtype_lower}.")

    for old_key, config in LEGACY_ARGS.items():
        new_key = config.get("new_key")
        if new_key in params_dict and params_dict[new_key] is None:
            params_dict[new_key] = params_dict.pop(old_key)
        else:
            params_dict.pop(old_key)
